#include "stream_reassembler.hh"

// Dummy implementation of a stream reassembler.

// For Lab 1, please replace with a real implementation that passes the
// automated checks run by `make check_lab1`.

// You will need to add private members to the class declaration in `stream_reassembler.hh`

template <typename... Targs>
void DUMMY_CODE(Targs &&.../* unused */) {}

using namespace std;

StreamReassembler::StreamReassembler(const size_t capacity) : _output(capacity), _capacity(capacity) {}

//! \details This function accepts a substring (aka a segment) of bytes,
//! possibly out-of-order, from the logical stream, and assembles any newly
//! contiguous substrings and writes them into the output stream in order.
void StreamReassembler::push_substring(const string &data, const size_t index, const bool eof) {
  if (eof) {
    _eof = true;
    _eof_idx = index + data.size();
  }
  // 先查看有没有可以及时push的数据
  // if (!_memo.empty()) {
  //     auto iter = _memo.begin();
  //     if (iter->first == _next_assembled_idx) {
  //         auto& content = iter->second;
  //         auto write_count = _output.write(content);
  //         _next_assembled_idx += write_count;
  //         _unassembled_bytes -= content.size();
  //         _memo.erase(iter);
  //         if (write_count < content.size()) {
  //             _unassembled_bytes += content.size() - write_count;
  //             _memo.emplace(_next_assembled_idx, content.substr(write_count));
  //         }
  //     }
  // }
  size_t idx = index;
  if (idx <= _next_assembled_idx) {
    if (idx + data.size() > _next_assembled_idx) {
      // 未写完的直接丢弃
      auto content = data.substr(_next_assembled_idx - idx);
      auto write_count = _output.write(content);
      _next_assembled_idx += write_count;
      // 优化丢弃
      // if (write_count < content.size()) {
      //     auto s = content.substr(write_count);
      //     // 能否和右边合并
      //     auto iter = _memo.begin();
      //     while (iter != _memo.end() && iter->first <= _next_assembled_idx + s.size()) {
      //         auto sz = iter->second.size();
      //         if (iter->first + sz > _next_assembled_idx + s.size()) {
      //             s += std::move(iter->second).substr(_next_assembled_idx + s.size() - iter->first);
      //         }
      //         _unassembled_bytes -= sz;
      //         iter = _memo.erase(iter);
      //     }
      //     auto store_count = std::min(s.size(), _capacity - _output.buffer_size() - _unassembled_bytes);
      //     _unassembled_bytes += store_count;
      //     _memo.emplace(_next_assembled_idx, s.substr(0, store_count));
      // }

      auto iter = _memo.begin();
      while (iter != _memo.end() && iter->first <= _next_assembled_idx) {
        auto sz = iter->second.size();
        if (iter->first + sz >= _next_assembled_idx) {
          // 未写完的直接丢弃, 哪怕是已经保存的数据
          auto existed_str = iter->second.substr(_next_assembled_idx - iter->first);
          auto add_count = _output.write(existed_str);
          _next_assembled_idx += add_count;
          // if (add_count < existed_str.size()) {
          //     _memo.erase(iter);
          //     _unassembled_bytes -= existed_str.size();
          //     _unassembled_bytes += existed_str.size() - write_count_2;
          //     _memo.emplace(_next_assembled_idx, existed_str.substr(write_count_2));
          //     break;
          // }
        }
        _unassembled_bytes -= sz;
        iter = _memo.erase(iter);
      }
    }
  } else {
    string s = data;
    auto iter = _memo.lower_bound(idx);
    // 后向合并
    while (iter != _memo.end() && iter->first <= idx + s.size()) {
      auto sz = iter->second.size();
      if (iter->first + sz > idx + s.size()) {
        s += iter->second.substr(idx + s.size() - iter->first);
      }
      _unassembled_bytes -= sz;
      iter = _memo.erase(iter);
    }
    // 前向合并
    iter = _memo.lower_bound(idx);
    if (iter != _memo.begin() && !_memo.empty()) {
      iter--;
      auto sz = iter->second.size();
      // 检查是否可以和前一个合并
      if (iter->first + sz >= idx) {
        if (iter->first + sz > idx + s.size()) {
          s = iter->second;
        } else {
          s = iter->second + s.substr(iter->first + sz - idx);
        }
        _unassembled_bytes -= sz;
        idx = iter->first;
        _memo.erase(iter);
      }
    }
    auto store_count = std::min(s.size(), _capacity - _output.buffer_size() - _unassembled_bytes);
    _unassembled_bytes += store_count;
    _memo.emplace(idx, s.substr(0, store_count));
  }
  if (_eof && _next_assembled_idx >= _eof_idx) {
    _output.end_input();
  }
}

size_t StreamReassembler::unassembled_bytes() const { return _unassembled_bytes; }

bool StreamReassembler::empty() const { return _unassembled_bytes == 0; }
